import asyncio
import time
import uuid
from typing import Optional

import redis.asyncio as redis


class AsyncRedisLock:
    """基于官方redis-py库的异步分布式锁"""

    def __init__(
        self,
        redis_client: redis.Redis,
        lock_key: str,
        lock_prefix_key: str = None,
        timeout: int = 30,
        retry_delay: float = 0.1,
    ):
        """
        初始化异步Redis分布式锁

        Args:
            redis_client: Redis客户端实例
            lock_key: 锁的键名
            lock_prefix_key: 锁的前缀键名（可选）
            timeout: 锁的超时时间（秒）
            retry_delay: 重试间隔（秒）
        """
        self.redis = redis_client
        __lock_prefix_key = lock_prefix_key or "distributed_lock"
        self.lock_key = f"{__lock_prefix_key}:{lock_key}"
        self.timeout = timeout
        self.retry_delay = retry_delay
        self.identifier = str(uuid.uuid4())
        self._locked = False

    async def acquire(self, blocking: bool = True, timeout: Optional[float] = None) -> bool:
        """
        获取锁

        Args:
            blocking: 是否阻塞等待
            timeout: 阻塞等待的超时时间

        Returns:
            bool: 是否成功获取锁
        """
        end_time = None
        if timeout is not None:
            end_time = time.time() + timeout

        while True:
            # 使用SET命令的NX和EX参数实现原子操作
            result = await self.redis.set(
                self.lock_key, self.identifier, nx=True, ex=self.timeout  # 只在键不存在时设置  # 设置过期时间
            )

            if result:
                self._locked = True
                return True

            if not blocking:
                return False

            # 检查是否超时
            if end_time and time.time() > end_time:
                return False

            # 等待后重试
            await asyncio.sleep(self.retry_delay)

    async def release(self) -> bool:
        """
        释放锁

        Returns:
            bool: 是否成功释放锁
        """
        if not self._locked:
            return False

        # 使用Lua脚本确保原子性：只有锁的持有者才能释放锁
        lua_script = """
        if redis.call("GET", KEYS[1]) == ARGV[1] then
            return redis.call("DEL", KEYS[1])
        else
            return 0
        end
        """

        try:
            result = await self.redis.eval(lua_script, 1, self.lock_key, self.identifier)
            self._locked = False
            return bool(result)
        except Exception as e:
            # 发生异常时也标记为未锁定状态
            self._locked = False
            return False

    async def extend(self, additional_time: int) -> bool:
        """
        延长锁的过期时间

        Args:
            additional_time: 延长时间（秒）

        Returns:
            bool: 是否成功延长
        """
        if not self._locked:
            return False

        lua_script = """
        if redis.call("GET", KEYS[1]) == ARGV[1] then
            return redis.call("EXPIRE", KEYS[1], ARGV[2])
        else
            return 0
        end
        """

        try:
            result = await self.redis.eval(lua_script, 1, self.lock_key, self.identifier, str(additional_time))
            return bool(result)
        except Exception:
            return False

    async def __aenter__(self):
        """异步上下文管理器入口"""
        acquired = await self.acquire()
        if not acquired:
            raise RuntimeError("Failed to acquire lock")
        return self

    async def __aexit__(self, exc_type, exc_val, exc_tb):
        """异步上下文管理器出口"""
        await self.release()
        return False

    @property
    def locked(self) -> bool:
        """检查锁是否被当前实例持有"""
        return self._locked


class AsyncLockManager:
    """异步锁管理器"""

    def __init__(self, redis_url: str = "redis://localhost:6379/0"):
        self.redis_url = redis_url
        self._redis_client = None

    async def get_redis_client(self) -> redis.Redis:
        """获取Redis客户端"""
        if self._redis_client is None:
            self._redis_client = redis.Redis.from_url(self.redis_url)
        return self._redis_client

    async def get_lock(self, lock_key: str, timeout: int = 30, retry_delay: float = 0.1) -> AsyncRedisLock:
        """获取锁实例"""
        redis_client = await self.get_redis_client()
        return AsyncRedisLock(redis_client, lock_key, timeout, retry_delay)

    async def close(self):
        """关闭Redis连接"""
        if self._redis_client:
            await self._redis_client.aclose()


# 使用示例
async def example_basic_usage():
    """基础使用示例"""
    # 方式1：直接创建Redis客户端
    redis_client = redis.Redis.from_url("redis://localhost:6379/0")

    try:
        # 创建锁实例
        lock = AsyncRedisLock(redis_client, "my_resource", timeout=10)

        # 方式1：手动获取和释放
        if await lock.acquire():
            try:
                print("获取锁成功，执行业务逻辑...")
                await asyncio.sleep(2)
                print("业务逻辑执行完成")
            finally:
                await lock.release()
                print("锁已释放")
        else:
            print("获取锁失败")

        # 方式2：使用异步上下文管理器
        try:
            async with AsyncRedisLock(redis_client, "another_resource", timeout=5) as lock:
                print("在锁保护下执行业务逻辑...")
                await asyncio.sleep(1)
                print("业务逻辑执行完成")
        except RuntimeError as e:
            print(f"获取锁失败: {e}")

    finally:
        await redis_client.aclose()


async def example_concurrent_access():
    """并发访问示例"""
    lock_manager = AsyncLockManager()

    async def worker(worker_id: int):
        """工作协程"""
        lock = await lock_manager.get_lock("shared_resource", timeout=5)

        # 尝试非阻塞获取锁
        if await lock.acquire(blocking=False):
            try:
                print(f"Worker {worker_id} 获取锁成功")
                # 模拟业务处理
                await asyncio.sleep(2)
                print(f"Worker {worker_id} 业务处理完成")
            finally:
                await lock.release()
                print(f"Worker {worker_id} 释放锁")
        else:
            print(f"Worker {worker_id} 获取锁失败，跳过任务")

    # 启动多个并发任务
    tasks = [worker(i) for i in range(3)]
    await asyncio.gather(*tasks)

    await lock_manager.close()


async def example_blocking_wait():
    """阻塞等待示例"""
    lock_manager = AsyncLockManager()

    async def slow_worker():
        """慢速工作协程"""
        lock = await lock_manager.get_lock("blocking_resource", timeout=10)
        async with lock:
            print("慢速工作协程获取锁，开始长时间工作...")
            await asyncio.sleep(3)
            print("慢速工作协程完成")

    async def fast_worker():
        """快速工作协程"""
        lock = await lock_manager.get_lock("blocking_resource", timeout=10)
        # 阻塞等待获取锁，最多等待5秒
        if await lock.acquire(blocking=True, timeout=5.0):
            try:
                print("快速工作协程获取锁")
                await asyncio.sleep(1)
                print("快速工作协程完成")
            finally:
                await lock.release()
        else:
            print("快速工作协程等待超时")

    # 启动任务
    slow_task = asyncio.create_task(slow_worker())
    await asyncio.sleep(0.1)  # 确保慢速任务先获取锁
    fast_task = asyncio.create_task(fast_worker())

    await asyncio.gather(slow_task, fast_task)

    await lock_manager.close()


async def example_lock_extension():
    """锁延长效期示例"""
    redis_client = redis.Redis.from_url("redis://localhost:6379/0")

    try:
        lock = AsyncRedisLock(redis_client, "extendable_resource", timeout=3)

        async with lock:
            print("获取锁，初始超时3秒")
            await asyncio.sleep(2)
            print("2秒后，延长锁5秒")

            if await lock.extend(5):
                print("锁延长效期成功")
                await asyncio.sleep(3)
                print("再等待3秒，锁仍然有效")
            else:
                print("锁延长效期失败")

            print("业务逻辑完成")

    finally:
        await redis_client.aclose()


# 安装依赖
"""
pip install redis
"""

if __name__ == "__main__":
    print("=== 基础使用示例 ===")
    asyncio.run(example_basic_usage())

    print("\n=== 并发访问示例 ===")
    asyncio.run(example_concurrent_access())

    print("\n=== 阻塞等待示例 ===")
    asyncio.run(example_blocking_wait())

    print("\n=== 锁延长效期示例 ===")
    asyncio.run(example_lock_extension())
